import numpy as np


class CosineAnnealing:
    def __init__(self, lr_min, lr_max, total_iters,
                 warmup_iters=0, lr_init=None):
        self.lr_min = lr_min
        self.lr_max = lr_max
        self.warmup_iters = warmup_iters
        self.total_iters = total_iters - self.warmup_iters
        self.lr_init = lr_init

    def __call__(self, optimizer, cur_iter):
        if cur_iter < self.warmup_iters:
            lr = self.lr_init + cur_iter * (self.lr_max - self.lr_init) / self.warmup_iters
        else:
            lr = self.lr_min + (self.lr_max - self.lr_min) * (1 + np.cos(cur_iter * np.pi / self.total_iters)) / 2
        optimizer.param_groups[0]["lr"] = lr

        return lr


class LinearDecay:
    def __init__(self, lr_min, lr_max, total_iters,
                 warmup_iters=0, lr_init=None):
        self.lr_min = lr_min
        self.lr_max = lr_max
        self.warmup_iters = warmup_iters
        self.total_iters = total_iters - self.warmup_iters
        self.lr_init = lr_init

    def __call__(self, optimizer, cur_iter):
        if cur_iter < self.warmup_iters:
            lr = self.lr_init + cur_iter * (self.lr_max - self.lr_init) / self.warmup_iters
        else:
            # lr = self.lr_min + (self.lr_max - self.lr_min) * (1 + np.cos(cur_iter * np.pi / self.total_iters)) / 2
            lr = self.lr_max + (cur_iter - self.warmup_iters) * (self.lr_min - self.lr_max) / self.total_iters
        optimizer.param_groups[0]["lr"] = lr

        return lr
